import os
import json
import requests
from tqdm import tqdm
import time

class DeepSeekAPI:
    def __init__(self, api_key=None):
        self.api_key = api_key or os.getenv('DEEPSEEK_API_KEY')
        self.api_url_root = os.getenv('DEEPSEEK_API_URL') if os.getenv('DEEPSEEK_API_URL') else "https://api.deepseek.com/v1"
        self.model = os.getenv('DEEPSEEK_MODEL') if os.getenv('DEEPSEEK_MODEL') else "deepseek-chat"
        if not self.api_key:
            raise ValueError("DeepSeek API key is required")
        
        self.api_url = self.api_url_root+"/chat/completions"  # 替换为实际的API端点
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
    
    def generate_questions(self, prompt):
        try:
            payload = {
                "model": self.model,
                "messages": [
                    {"role": "user", "content": prompt}
                ],
                "temperature": 0.7,
                "stream": True
            }
            
            response = requests.post(
                self.api_url,
                headers=self.headers,
                json=payload,
                stream=True
            )
            
            response.raise_for_status()
            
            collected_messages = []
            total_tokens = 0
            start_time = time.time()
            pbar = tqdm(desc="生成进度", unit="tokens", dynamic_ncols=True)
            
            for line in response.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith('data: '):
                        try:
                            json_data = json.loads(line[6:])
                            if json_data.get('choices'):
                                content = json_data['choices'][0].get('delta', {}).get('content', '')
                                if content:
                                    collected_messages.append(content)
                                    token_count = len(content.split())
                                    total_tokens += token_count
                                    pbar.update(token_count)
                                    
                                    elapsed_time = time.time() - start_time
                                    if elapsed_time > 0:
                                        tokens_per_second = total_tokens / elapsed_time
                                        progress = min(90, int((total_tokens / 500) * 100))
                                        yield {
                                            'progress': progress,
                                            'speed': f'{tokens_per_second:.1f}',
                                            'total_tokens': total_tokens,
                                            'elapsed_time': f'{elapsed_time:.1f}'
                                        }
                        except json.JSONDecodeError:
                            continue
            
            pbar.close()
            full_response = ''.join(collected_messages)
            
            if not full_response.strip():
                yield {"error": "Empty response received from API"}
                return
            
            yield {'progress': 95, 'speed': '0', 'status': '正在解析响应数据'}
            
            try:
                # print(full_response)
                questions_data = json.loads(full_response)

                yield {'progress': 100, 'speed': '0', 'status': '完成'}
                yield questions_data
            except json.JSONDecodeError:
                cleaned_response = full_response.strip()
                if cleaned_response.endswith(','):
                    cleaned_response = cleaned_response[:-1]
                try:
                    questions_data = json.loads(cleaned_response)
                    yield {'progress': 100, 'speed': '0', 'status': '完成'}
                    yield questions_data
                except json.JSONDecodeError:
                    yield {'progress': 100, 'speed': '0', 'error': '无法解析API返回的JSON数据，请重试'}
                
        except requests.exceptions.RequestException as e:
            return {"error": f"API 请求失败: {str(e)}"}
        except Exception as e:
            return {"error": f"发生错误: {str(e)}"}